製造ビジネステクノロジー部のやまたつです。

re:Invent 2024でのセッション「DAT402-R1 | Using LangChain to build gen AI apps with Amazon Aurora and pgvector [REPEAT]」に参加しましたので、その内容を共有いたします。
このセッションはCode Talk形式で行われ、LangChain, pgvector, Bedrock を使ってsemantic searchを実装する方法について、PythonコードとSQLを交えて説明が行われました。

英語を聞くのとコードを読むのとで必死になっていたため写真少なめです。。。🙇

セッションの流れ

まず、LangChain, pgvector, Bedrockについての簡単な説明があったと、早速Pythonコードが映し出されてLangChainを用いてアプリケーションが実装され始めました。
Pythonコード以外にも、PostgresのCLIコンソールを使ってデータベースを操作するのをライブで見ることができました。クエリ結果をJSONに変換して、サブクエリの結果を親JSONに埋め込んで、JSONをvectorに変換して、とスラスラと操作しているのが印象的でした。
（CLIでDBを操作しているときの全能感、ハンパないよね）

このセッションで作成したアプリケーションでは以下の機能が実装されていました。

1. ユーザーは映画に関する質問を自然言語で入力
2. 入力内容をLangChain, Bedrockを用いてembeddingに変換
3. embeddingを用いてデータベースに対して検索クエリを発行
4. 取得された映画情報をBedrockによって自然言語に変換し、ユーザーに返す

また、上記の機能を実現するために、以下の前処理が行われていました。

1. 架空の映画データベースを作成
2. 映画情報をPostgresのSQLを用いてJSON形式に変換し、新たなテーブルとしてPostgresに格納
3. JSON形式のデータをLangChainによってembeddingに変換
4. embeddingを新たなテーブルとしてPostgresに格納
5. embeddingにpgvectorを用いてインデックスを作成し、検索クエリの高速化を図る

以下では特に技術的なポイントについて詳しく説明していきます。

技術的なポイント

PostgresのクエリによるJSON形式データの生成

このセッションではPostgresのCLIコンソールによるデータ操作に時間を使っていました。
その中で、json_build_objectやjson_aggなどのPostgresの関数を使って、映画情報をJSON形式に変換していました。

正確なクエリは覚えていないのですが、以下のようにサブクエリを利用してactorテーブルの検索結果をmovieテーブルのデータに埋め込んでいました。

SELECT
  json_build_object(
    'title', title,
    'actors', (
      SELECT json_agg(
        json_build_object(
          'name', name,
          'age', age
        )
      )
      FROM actor
      WHERE movie_id = movie.id
    )
  )

JSON形式データのembeddingへの変換

SQLによって取得されたJSONデータはLangChainおよびBedrockによってembeddingに変換されていました。
（曖昧な記憶から再現しているので一部セッションの内容と異なる箇所があるかもしれません。。。🙇）

import json
from langchain.embeddings import BedrockEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores.pgvector import PGVector
from langchain.schema import Document
import boto3

 Bedrockクライアントの設定
bedrock_client = boto3.client(
    service_name="bedrock-runtime",
    region_name="ap-northeast-1"  # リージョンを指定
)

 Bedrockのembeddingsモデルを初期化
embeddings = BedrockEmbeddings(
    client=bedrock_client,
    model_id="amazon.titan-embed-text-v1"  # または他のモデル
)

 データベース接続情報
CONNECTION_STRING = "postgresql://user:password@host:5432/dbname"

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100,
    separators=["\n\n", "\n", " ", ""]
)

 JSONデータの取得
json_data = fetch_json_data()

 ドキュメントの準備
docs = []
for item in json_data:
    doc = Document(
        page_content=json.dumps(item),
        metadata={"id": item['id']}
    )
    docs.append(doc)

 PGVectorストアの設定と保存
collection_name = "documents_embeddings"
PGVector.from_documents(
    documents=text_splitter.split_documents(docs),
    embedding=embeddings,
    collection_name=collection_name,
    connection_string=CONNECTION_STRING,
    pre_delete_collection=True  # 既存のコレクションを削除する場合
)

fetch_json_data()の中で使っていたSQLは、PostgresのCLIコンソールでコネコネ作り上げていって、最後にコピペして完成させていました。

embeddingのテーブル作成とインデックスの設定

PostgresのCLIコンソールを使って、embeddingのテーブルを作成し、pgvectorのhnsw関数を使ってcosine類似度インデックスを設定していました。
（こちらも、記憶を基にAIと協力して復元しているため一部セッションの内容と異なる箇所があるかもしれません🙇🙇🙇）

CREATE TABLE documents_embeddings (
  id SERIAL PRIMARY KEY,
  embedding VECTOR(1024)
);

CREATE INDEX embedding_idx ON documents_embeddings USING hnsw(embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);

vector_halfを利用した改善

セッション中では、vector_halfを使って改善を行っていました。
初学者である筆者からすると簡単な修正によってクエリが高速化されていて、豆鉄砲を100発喰らった鳩になりました。鳩になりました。

DROP INDEX embedding_idx;

ALTER TABLE your_table
ALTER COLUMN embedding TYPE vector_half(1024);

CREATE INDEX ON items USING hnsw (embedding vector_half_cosine_ops) WITH (m = 16, ef_construction = 64);

セッションではこの変更によって以下のようにデータベースの全体的なパフォーマンスが向上が得られると話されていました。

• キャッシュに収まるデータ量が増え、クエリの実行速度が向上する
• テーブルの総容量が小さくなり、ディスク I/O が減少する
• 小さなデータ型を使うことで、メモリ使用量を削減できる

一方で、質疑も含めて、この変更によるデメリットについて触れられていました。

• 検索結果の一貫性が低下すること
  • 具体的には、同じクエリを複数回実行したにもかかわらず、異なる結果が返ってくることがある
• また、質問者から、ハイブリッド検索の手法を使うことで改善可能である点が指摘されていました。
  • 単純な semantic search だけでなく、他の検索アルゴリズムと組み合わせることで、より安定した検索結果が得られるとのことでした。

まとめ

今回のセッションはCode Talkというコードを交えた説明形式で行われていました。
LangChainもGen AIも初学者でようやくembeddingがわかった程度の筆者が参加するには濃厚すぎる会でしたが、コードやSQL操作と合わせて説明されることでかなり理解しやすかったです。

一緒に参加された日本語話者の方にvector_halfの件について質問していただいたのですが、その時はちゃんと回答できなくて申し訳なかったです。その後、transcribeの情報からこのブログ作成にこぎつけたので、どうか質問していただいた方に届いて欲しいです。どうか！！！

以上でした！